import subprocess
import re

mirrors = [
    "https://pypi.tuna.tsinghua.edu.cn/simple",
    "https://mirrors.aliyun.com/pypi/simple",
    "https://pypi.mirrors.ustc.edu.cn/simple"
]

def check_mirror_speed(mirror_url):
    try:
        # 测试镜像源响应时间
        result = subprocess.run(
            ["pip", "install", "--dry-run", "--quiet", "requests", "-i", mirror_url],
            capture_output=True,
            text=True,
            timeout=10
        )
        if "Could not find a version" in result.stderr:
            return float('inf')  # 镜像源不可用
        # 提取下载速度（示例输出解析，实际需根据 pip 版本调整）
        speed_match = re.search(r"(\d+\.\d+[KM]B/s)", result.stdout)
        speed = float(speed_match.group(1).replace("KB/s", "1").replace("MB/s", "1000")) if speed_match else 0
        return speed
    except:
        return float('inf')

# 按速度排序镜像源
sorted_mirrors = sorted(mirrors, key=lambda x: check_mirror_speed(x))

# 更新 pip 配置文件
with open("/Users/jxrt_1/.pip/pip.conf", "w") as f:
    f.write("[global]\n")
    f.write(f"index-url = {sorted_mirrors[0]}\n")
    f.write("trusted-host = ")
    f.write("\n    ".join([url.split("//")[1].split("/")[0] for url in mirrors]))
    print(f"已自动切换至最快镜像源: {sorted_mirrors[0]}")